#!/beegfs/group/lqcd/software/anaconda3/bin/python3
# encoding: utf-8


import numpy as np
from matplotlib import pyplot as plt
from math import exp, ceil, log
from tqdm import tqdm
import random as rand
import gvar as gv



N_t = 32
dt = 0.25
N_cnfg = 8000
binsz = 40
amp = 0.8
N_crrltn = ceil(1/dt**2)
N_thrmlzng = 4 * N_crrltn


def S_local(xlst, j):
    jp = (j+1) % N_t
    jm = (j-1) % N_t
    return dt*xlst[j]**2/2 + xlst[j]*(xlst[j]-xlst[jp]-xlst[jm])/dt


def Mtrpls_Hstng(xlst):
    global acptd
    for j in range(0, N_t):
        x_orgnl = xlst[j]
        S_local_orgnl = S_local(xlst, j)
        xlst[j] = xlst[j] + rand.uniform(-amp, amp)
        acptd = acptd + 1
        dS = S_local(xlst, j) - S_local_orgnl
        xi = rand.uniform(0, 1)
        if exp(-dS) < xi:
            xlst[j] = x_orgnl
            acptd = acptd - 1
    return xlst


acptd = 0
xlst = np.zeros(N_t)
print('>>> Thermalizing ......',)
for i in tqdm(range(0, N_thrmlzng)):
    xlst = Mtrpls_Hstng(list(xlst))
print('    Acceptance: ', acptd/(N_thrmlzng*N_t))

acptd = 0
cnfg_lst = []
print('>>> Gnerating configurations ......',)
for i in tqdm(range(0, N_cnfg)):
    for j in range(0, N_crrltn):
        xlst = Mtrpls_Hstng(list(xlst))
    cnfg_lst.append(list(xlst))
print('    Acceptance: ', acptd/(N_cnfg*N_t*N_crrltn))


cnfgs = np.array(cnfg_lst)

G_lst = np.array([1/N_t*np.sum(np.array([cnfgs[:,j]*cnfgs[:,(j+t) % N_t] for j in range(0, N_t)]),axis=0) for t in range(0,N_t)]).T

def bin(G,binsize):
    G_binned = []
    for i in range(0,len(G),binsize):
        G_binned.append(np.mean(G[i:i+binsize],axis=0))      
    return G_binned

G_lst_bnd = np.array(bin(G_lst,binsz))

c2pnt_cntrl = np.mean(G_lst_bnd,axis=0)
c2pnt_err = np.var(G_lst_bnd,axis=0)

c_2pnt = gv.gvar(c2pnt_cntrl,c2pnt_err)

log_r_lst = gv.log(c_2pnt[0:-1]/c_2pnt[1:])/dt
eng_cntrl = [log_r.mean for log_r in log_r_lst]
eng_err = [log_r.sdev for log_r in log_r_lst]


fig = plt.figure(figsize=(20,6))

t_lst=np.linspace(0,N_t-1,N_t)
c2pnt_plt=fig.add_subplot(121,xlim=(0.2,N_t),ylim=(-0.075,1.05*max(c2pnt_cntrl)))
c2pnt_plt.errorbar(t_lst,c2pnt_cntrl,yerr=c2pnt_err,fmt='.',ecolor='r',color='b',elinewidth=2,capsize=4,markersize='10')
c2pnt_plt.set_title('$C_2(t)$',fontsize=18)



t_lst=np.linspace(0,N_t-2,N_t-1)
eng_plt=fig.add_subplot(122,xlim=(0,int(N_t)/3),ylim=(0,2))  
eng_plt.errorbar(t_lst,eng_cntrl,yerr=eng_err,fmt='.',ecolor='r',color='b',elinewidth=2,capsize=4,markersize='10')
eng_plt.set_title('$\Delta E_{eff}(t)$',fontsize=18)

plt.savefig('MCMC_hrmnc_oscltr.png',dpi=400)